#include "defs.h"
#include <stddef.h> 

pagetable_t kernel_pagetable; //存储页表内存地址
extern char etext[], _start[], _etext[];  
static void handle_page_fault(uint64 fault_va); //页故障处理函数


//创建页表
pagetable_t kvmmake(void) {
  pagetable_t kpgtbl = (pagetable_t)alloc_page();
  if (kpgtbl == NULL) {
    panic("kvmmake: failed to allocate root pagetable");
  }
  memset(kpgtbl, 0, PGSIZE);

  // 映射内核代码段（确保大小页对齐）
  uint64 kernel_text_start = (uint64)_start;  // 内核代码起始地址
  uint64 kernel_text_end = (uint64)_etext;    // 内核代码结束地址
  
  // 计算大小并确保页对齐
  uint64 text_size = kernel_text_end - kernel_text_start;
  text_size = PGROUNDUP(text_size);  // 向上对齐到页大小
  
  // 检查范围有效性
  if (kernel_text_start + text_size > PHYSTOP) {
    panic("kvmmake: kernel text exceeds PHYSTOP");
  }
  
  // 映射内核代码段，权限为可读可执行
  if (mappages(kpgtbl, kernel_text_start, text_size, 
               kernel_text_start, PTE_R | PTE_X | PTE_V) != 0) {
    panic("kvmmake: failed to map kernel text");
  }

  // 映射UART设备（调试输出必需）
  if (mappages(kpgtbl, UART0, PGSIZE, UART0, PTE_R | PTE_W | PTE_V) != 0) {
    panic("kvmmake: failed to map UART");
  }

  return kpgtbl;
}

void kvmmap(pagetable_t kpgtbl, uint64 va, uint64 pa, uint64 sz, int perm) {
  if (mappages(kpgtbl, va, sz, pa, perm) != 0)
    panic("kvmmap");
}

void kvminit(void) {
  kernel_pagetable = kvmmake();
  if (kernel_pagetable == NULL) {
    panic("kvminit: kernel_pagetable is NULL");
  }
  
  // 注册异常处理入口：设置stvec寄存器指向trap_handler
  w_stvec((uint64)trap_handler);
}

void kvminithart() {
  // 刷新TLB并启用分页
  sfence_vma();
  w_satp(MAKE_SATP(kernel_pagetable));
  sfence_vma();
}

pte_t *walk(pagetable_t pagetable, uint64 va, int alloc) {
  if (va >= MAXVA)
    panic("walk");

  for (int level = 2; level > 0; level--) {
    pte_t *pte = &pagetable[PX(level, va)];
    if (*pte & PTE_V) {
      pagetable = (pagetable_t)PTE2PA(*pte);
    } else {
      if (!alloc || (pagetable = (pagetable_t)alloc_page()) == 0)
        return 0;
      memset(pagetable, 0, PGSIZE);
      *pte = PA2PTE(pagetable) | PTE_V;
    }
  }
  return &pagetable[PX(0, va)];
}

int mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm) {
  uint64 a, last;
  pte_t *pte;

  // 检查参数合法性
  if (va % PGSIZE != 0) {
    printf("mappages: va not aligned - 0x%lx\n", va);
    return -1;
  }
  if (pa % PGSIZE != 0) {
    printf("mappages: pa not aligned - 0x%lx\n", pa);
    return -1;
  }
  if (size == 0) {
    printf("mappages: size is zero\n");
    return -1;
  }
  if (size % PGSIZE != 0) {
    printf("mappages: size not page-aligned - 0x%lx\n", size);
    return -1;
  }
  if (va + size > MAXVA) {
    printf("mappages: va out of range - 0x%lx\n", va + size);
    return -1;
  }

  a = va;
  last = va + size - PGSIZE;
  for (;;) {
    if ((pte = walk(pagetable, a, 1)) == 0) {
      printf("mappages: walk failed for va 0x%lx\n", a);
      return -1;
    }
    if (*pte & PTE_V) {
      printf("mappages: remap of va 0x%lx\n", a);
      return -1;
    }
    *pte = PA2PTE(pa) | perm | PTE_V;
    if (a == last)
      break;
    a += PGSIZE;
    pa += PGSIZE;
  }
  return 0;
}

uint64 walkaddr(pagetable_t pagetable, uint64 va) {
  pte_t *pte;
  if (va >= MAXVA)
    return 0;
  pte = walk(pagetable, va, 0);
  if (pte == 0 || (*pte & PTE_V) == 0)
    return 0;
  return PTE2PA(*pte);
}

// 递归销毁某一级页表（从level=2到level=0）
static void free_single_pt(pagetable_t pt, int level) {
  // 三级页表（level=0）是叶子节点，没有子页表，直接释放当前页表
  if (level == 0) {
    free_page(pt);
    return;
  }

  // 遍历当前页表的所有PTE（共512个）
  for (int i = 0; i < 512; i++) {
    pte_t pte = pt[i];
    // 检查PTE是否有效，且指向的是子页表（非物理页）
    if (pte & PTE_V) {
      // 从PTE中提取子页表的物理地址（PTE2PA：去掉低10位标志位）
      pagetable_t child_pt = (pagetable_t)PTE2PA(pte);
      // 递归销毁子页表（级别-1：比如level=2→level=1）
      free_single_pt(child_pt, level - 1);
      // 清除当前PTE（标记为无效，避免后续误访问）
      pt[i] = 0;
    }
  }

  // 释放当前级页表的物理页
  free_page(pt);
}

// 销毁根页表
void destroy_pagetable(pagetable_t pt) {
  // 检查页表是否为空，避免空指针访问
  if (pt == NULL) {
    panic("destroy_pagetable: null pagetable");
  }
  // 从根页表（level=2）开始递归销毁
  free_single_pt(pt, 2);
}
    
void handle_page_fault(uint64 fault_va) {
  printf("Page fault at virtual address: 0x%lx\n", fault_va);
  panic("Unhandled page fault"); // 基础处理：打印并终止
}

// 异常处理入口
void trap_handler() {
  uint64 scause = r_scause();  // 获取异常原因

  // 判断是否为页故障（13=指令页故障，15=数据页故障）
  if (scause == 13 || scause == 15) {
    uint64 fault_va = r_stval();  // 获取故障虚拟地址
    handle_page_fault(fault_va);
  } else {
    // 其他异常（如中断），简单 panic
    panic("Unexpected trap: scause=0x%lx", scause);
  }
}

